local Node = {}
Node.__cname = "util.BRTree.Node"
Node.__index = Node

function Node.new(v)
    local obj = {}
    setmetatable(obj, Node)
    obj:ctor(v)
    return obj
end

function Node:ctor(v)
    self.value = v
    self.left = nil
    self.right = nil
    self.isBlack = false
end

local BRTree = {}
BRTree.__index = BRTree

local BRTree = {}
BRTree.__cname = "util.BRTree"
BRTree.__index = BRTree

function BRTree.new(cmpFunc)
    local obj = {}
    setmetatable(obj, BRTree)
    obj:ctor(cmpFunc)
    return obj
end

function BRTree:ctor(cmpFunc)
    self.cmpFunc = cmpFunc
    self.count = 0
    self.root = nil
end

function BRTree:Clear()
    self.count = 0
    self.root = nil
end

function BRTree:GetCount()
    return self.count
end

function BRTree:Add(v)
    local newNode = Node.new(v)
    if nil == self.root then --情况1
        self.root = newNode
        newNode.isBlack = true
        self.count = self.count + 1
        return
    end
    --左节点比父节点小, 右节点比父节点大
    local stack = {}
    local node = self.root
    while nil ~= node do
        local cmpResult = self.cmpFunc(v, node.value)
        if cmpResult < 0 then -- 左子树
            table.insert(stack, node)
            if nil == node.left then
                node.left = newNode
                self:_AddFix(newNode, true, node, #stack, stack)
                self.count = self.count + 1
                break
            else
                node = node.left
            end
        elseif cmpResult > 0 then -- 右子树
            table.insert(stack, node)
            if nil == node.right then
                node.right = newNode
                self:_AddFix(newNode, false, node, #stack, stack)
                self.count = self.count + 1
                break
            else
                node = node.right
            end
        else
            print(tostring(v), "have exist ignore")
        end
    end
end


function BRTree:_AddFix(curNode, curNodeIsLeft, parent, parentLevel, stack)
    if self._noBlackRed then return end
    while nil ~= parent and not parent.isBlack do
        local gparent = stack[parentLevel-1]
        local uncle = nil
        local parentIsLeft = gparent.left == parent
        if parentIsLeft then
            uncle = gparent.right
        else
            uncle = gparent.left
        end

        if nil ~= uncle and not uncle.isBlack then --uncle存在且为红
            parent.isBlack = true
            uncle.isBlack = true
            if gparent == self.root then --gparent为root
                return
            end
            gparent.isBlack = false

            local ggparentLevel = parentLevel - 2
            local ggparent = stack[ggparentLevel]
            if ggparent.isBlack then
                return
            end
            --ggparent肯定存在且不为root
            curNode = gparent
            parentLevel = ggparentLevel
            parent = ggparent
            curNodeIsLeft = (parent.left == curNode)
        elseif nil == uncle or uncle.isBlack then --uncle不存在或为黑
            local ggparent = stack[parentLevel-2]
            if parentIsLeft then
                if curNodeIsLeft then --同侧1, 情况2-1
                    self:_RightRotate(gparent, ggparent)
                    parent.isBlack = true
                    gparent.isBlack = false
                else --不同侧1, 情况3-1
                    self:_LeftRotate(parent, gparent)
                    local temp = parent
                    parent = curNode
                    curNode = temp
                    self:_RightRotate(gparent, ggparent)
                    parent.isBlack = true
                    gparent.isBlack = false
                end
            else
                if curNodeIsLeft then --不同侧2, 情况3-2
                    self:_RightRotate(parent, gparent)
                    local temp = parent
                    parent = curNode
                    curNode = temp
                    self:_LeftRotate(gparent, ggparent)
                    parent.isBlack = true
                    gparent.isBlack = false
                else --同侧2, 情况2-2
                    self:_LeftRotate(gparent, ggparent)
                    parent.isBlack = true
                    gparent.isBlack = false
                end
            end
            return
        end
    end
end

function BRTree:_LeftRotate(subTreeRoot, subTreeParent)
    --print("left rotate")
    local newSubTreeRoot = subTreeRoot.right
    subTreeRoot.right = newSubTreeRoot.left
    newSubTreeRoot.left = subTreeRoot

    if subTreeRoot == self.root then
        self.root = newSubTreeRoot
    else
        if subTreeRoot == subTreeParent.left then
            subTreeParent.left = newSubTreeRoot
        else
            subTreeParent.right = newSubTreeRoot
        end
    end
end

---@param subTreeRoot "旋转节点"
---@param subTreeParent "旋转节点的父节点"
function BRTree:_RightRotate(subTreeRoot, subTreeParent)
    --print("right rotate")
    local newSubTreeRoot = subTreeRoot.left
    subTreeRoot.left = newSubTreeRoot.right
    newSubTreeRoot.right = subTreeRoot

    if subTreeRoot == self.root then
        self.root = newSubTreeRoot
    else
        if subTreeRoot == subTreeParent.left then
            subTreeParent.left = newSubTreeRoot
        else
            subTreeParent.right = newSubTreeRoot
        end
    end
end

function BRTree:_Find(v)
    local node = self.root
    while nil ~= node do
        local cmpResult = self.cmpFunc(v, node.value)
        if cmpResult < 0 then --比父节点小: 左侧
            node = node.left
        elseif cmpResult > 0 then --比父节点大: 右侧
            node = node.right
        else
            return node
        end
    end
    return nil
end

function BRTree:_SetColor(v, isBlack)
    if self.count < 1 then return end
    local node = self:_Find(v)
    if node then
        node.isBlack =isBlack
    end
end

function BRTree:Contains(v)
    if self.count < 1 then return false end
    return nil ~= self:_Find(v)
end

function BRTree:Remove(v)
    if self.count < 1 then return false end

    local stack = {}
    local parent = nil
    local node = self.root
    while nil ~= node do
        local cmpResult = self.cmpFunc(v, node.value)
        if cmpResult < 0 then --比父节点小: 左侧
            table.insert(stack, node)
            parent = node
            node = node.left
        elseif cmpResult > 0 then --比父节点大: 右侧
            table.insert(stack, node)
            parent = node
            node = node.right
        else
            self:_DelNode(node, parent, stack)
            self.count = self.count - 1
            return true
        end
    end
    return false
end

function BRTree:_DelNode(delNode, delNodeParent, stack)
    if nil == delNode.left and nil == delNode.right then --删除节点没有子节点, 直接删除(断开与parent的关系)
        self:_DelFix(delNode, stack)
        if delNode == self.root then
            self.root = nil
        elseif delNodeParent.left == delNode then
            delNodeParent.left = nil
        elseif delNodeParent.right == delNode then
            delNodeParent.right = nil
        else
            error("DelNode1")
        end
    elseif nil == delNode.left then --删除节点只有右子节点, 断开与parent的关系, parent子节点直接指向那个右子节点
        self:_DelFix(delNode, stack)
        if delNode == self.root then
            self.root = delNode.right
        elseif delNodeParent.left == delNode then
            delNodeParent.left = delNode.right
        elseif delNodeParent.right == delNode then
            delNodeParent.right = delNode.right
        else
            error("DelNode2")
        end
        delNode.right = nil --断开节点关系
    elseif nil == delNode.right then --删除节点只有左子节点, 断开与parent的关系, parent子节点直接指向那个左子节点
        self:_DelFix(delNode, stack)
        if delNode == self.root then
            self.root = delNode.left
        elseif delNodeParent.left == delNode then
            delNodeParent.left = delNode.left
        elseif delNodeParent.right == delNode then
            delNodeParent.right = delNode.left
        else
            error("DelNode3")
        end
        delNode.left = nil --断开节点关系
    else --左右子节点都有, 右子树中找一个最小的节点替换到删除节点处
        local replaceNode, replaceNodeParent = self:_FindReplaceNode(delNode, stack)
        delNode.value = replaceNode.value
        self:_DelNode(replaceNode, replaceNodeParent, stack)
    end
end

function BRTree:_FindReplaceNode(delNode, stack)
    local parent = delNode
    table.insert(stack, parent)
    local node = delNode.right
    while nil ~= node do
        if nil ~= node.left then
            table.insert(stack, node)
            parent = node
            node = node.left
        else
            return node, parent
        end
    end
    error("_FindReplaceNode")
end


---二叉搜索树的删除, 最终删除的节点: 叶子节点或没有左孩子的节点
function BRTree:_DelFix(delNode, stack)
    if self._noBlackRed then return end
    local parentLevel = #stack

    while true do
        if not delNode.isBlack then return end --删除节点为红色

        --删除节点为黑色
        if nil == delNode.left then
            if nil ~= delNode.right and not delNode.right.isBlack then --删除节点为黑, 只有右子节点, 且为红
                delNode.right.isBlack = true
                return
            end
        elseif nil == delNode.right then
            if nil ~= delNode.left and not delNode.left.isBlack then --删除节点为黑, 只有左子节点, 且为红
                delNode.left.isBlack = true
                return
            end
        end
        if delNode == self.root then return end

        --没有子节点, 子节点: 1黑1红, 2个黑, 2个红, 需要找sibling
        local parent = stack[parentLevel]
        local sibling = nil
        local siblingIsLeft = (parent.right == delNode)
        if siblingIsLeft then
            sibling = parent.left
        else
            sibling = parent.right
        end
        if sibling.isBlack then
            local siblingNoRedChild = true
            local farSiblingChildIsRed = false
            if nil == sibling.left then
                if nil ~= sibling.right then --1个右子节点, 肯定为红
                    siblingNoRedChild = false
                    if siblingIsLeft then --情况3-2, 近侄子为红
                        self:_LeftRotate(sibling, parent)
                        sibling.isBlack = false
                        sibling.right.isBlack = true
                        if 1 == parentLevel then return end
                        --情况4或5
                        parent.isBlack = true
                        sibling.isBlack = false
                        parentLevel = parentLevel - 1
                        delNode = parent
                    else
                        farSiblingChildIsRed = true
                    end
                end
            elseif nil == sibling.right then --1个左子节点, 肯定为红
                siblingNoRedChild = false
                if siblingIsLeft then
                    farSiblingChildIsRed = true
                else --情况3-1, 近侄子为红
                    local siblingLeft = sibling.left
                    self:_RightRotate(sibling, parent)
                    --旋转后变为情况4
                    sibling.isBlack = false
                    siblingLeft.isBlack = true
                end
            else --两个红或2个黑
                if siblingIsLeft then
                    farSiblingChildIsRed = not sibling.left.isBlack
                else
                    farSiblingChildIsRed = not sibling.right.isBlack
                end
                siblingNoRedChild = not farSiblingChildIsRed --要黑都是黑
            end

            if farSiblingChildIsRed then
                local gparent = stack[parentLevel-1]
                if siblingIsLeft then
                    self:_RightRotate(parent, gparent)
                    sibling.left.isBlack = true
                    sibling.isBlack = parent.isBlack
                    parent.isBlack = true
                else --情况2-1
                    self:_LeftRotate(parent, gparent)
                    sibling.right.isBlack = true
                    sibling.isBlack = parent.isBlack
                    parent.isBlack = true
                end
                return
            elseif siblingNoRedChild then
                if parent.isBlack then --情况5
                    sibling.isBlack = false
                    if 1 == parentLevel then return end
                    --继续往上处理
                    parentLevel = parentLevel - 1
                    delNode = parent
                else --情况4
                    parent.isBlack = true
                    sibling.isBlack = false
                    return
                end
            end
        else
            local gparent = stack[parentLevel - 1]
            if siblingIsLeft then --情况1-2
                local siblingRight = sibling.right
                self:_RightRotate(parent, gparent)
                --parent.isBlack = false
                sibling.isBlack = true
                --情况4
                siblingRight.isBlack = false
            else --情况1-1
                local siblingLeft = sibling.left
                self:_LeftRotate(parent, gparent)
                --parent.isBlack = false
                sibling.isBlack = true
                --情况4
                siblingLeft.isBlack = false
            end
            return
        end
    end
end

function BRTree:__tostring()
    if self.count < 1 then return "" end
    local strTb = {}
    local queue = {}
    table.insert(queue, self.root)

    while #queue > 0 do
        local node = queue[1]
        table.remove(queue, 1)
        if node.isBlack then
            table.insert(strTb, string.format("(%s)", tostring(node.value)))
        else
            table.insert(strTb, tostring(node.value))
        end

        if nil ~= node.left then
            table.insert(queue, node.left)
        end
        if nil ~= node.right then
            table.insert(queue, node.right)
        end
    end
    return table.concat(strTb, ",")
end



return BRTree

